"""
数据集操作方法
"""
import os
import cv2
import numpy as np
from video_classification import config


OUTPUT_LIST = config.OUTPUT_LIST
DATASET_DIR = 'datasets/videos/'


# 将数据集分割为训练数据集，测试数据集和交叉验证集三部分
def split_data(data_dir, output_list, ratio=0.7):
    """
    :param data_dir: 原始视频数据集的路径
    :param output_list: 获取原始视频的标签及其路径，将其分割为训练数据集，测试数据集和交叉验证数据集三部分，此为输出文件的路径
    :param ratio: 数据集分割比例
    :return: 无
    """
    # 读取所有视频的路径
    data_file = os.listdir(data_dir)
    data_list = []  # 存放视频路径
    label_list = []     # 存放视频标签
    for labels in data_file:
        vid_file = os.listdir(data_dir + labels)
        for vid in vid_file:
            vid_dir = os.path.join(data_dir, labels, vid)
            data_list.append(vid_dir)
            label_list.append(labels)
    permutation = np.random.permutation(len(data_list))     # 获取一个乱序的序列
    # 将打乱后的路径写入文件
    with open(output_list[0] + '.txt', 'w') as train:
        for data in permutation[0: int(ratio * len(data_list))]:
            train.write(data_list[data] + '\n')
    with open(output_list[1] + '.txt', 'w') as test:
        for data in permutation[int(ratio * len(data_list)): int((1 + ratio) / 2 * len(data_list))]:
            test.write(data_list[data] + '\n')
    with open(output_list[2] + '.txt', 'w') as val:
        for data in permutation[int((1 + ratio) / 2 * len(data_list)): len(data_list)]:
            val.write(data_list[data] + '\n')
    # 将打乱后的标签写入文件
    with open(output_list[0] + '_label.txt', 'w') as train:
        for label in permutation[0: int(ratio * len(label_list))]:
            train.write(label_list[label] + '\n')
    with open(output_list[1] + '_label.txt', 'w') as test:
        for label in permutation[int(ratio * len(label_list)): int((1 + ratio) / 2 * len(label_list))]:
            test.write(label_list[label] + '\n')
    with open(output_list[2] + '_label.txt', 'w') as val:
        for label in permutation[int((1 + ratio) / 2 * len(label_list)): len(label_list)]:
            val.write(label_list[label] + '\n')


# 从txt中读取数据集
def read_data(data_file, size):
    """
    :param data_file: 文件路径
    :param size: 输出视频尺寸
    :return: 字典类型，第一项为float32型的五维图像数组，第二项为int32型的二维标签数组，第三项为分类数
    """
    data_list = []
    label_list = []
    classes = []
    train_data = open(data_file + '.txt', 'rU').readlines()
    train_label = open(data_file + '_label.txt', 'rU').readlines()
    print(train_data)
    for i in range(len(train_data)):
        line_data = train_data[i].strip('\n')   # 取出行末的换行符
        line_label = train_label[i].strip('\n')
        if line_label not in classes:
            classes.append(line_label)
        # 获取视频路径
        vid = cv2.VideoCapture(line_data)
        # 统计帧率，不能使用OpenCV直接返回的值，不是真实值
        frame_number = 0
        while True:
            ret, frame = vid.read()
            frame_number += 1
            # 判断视频流中的帧是否存在
            if frame is None:
                break
        # 计算帧间隔
        span = frame_number / size[0]
        frame_span = []
        for j in range(size[0]):
            frame_span.append(int(span * j))
        frame_span.append(int(frame_number - 1))
        # 逐帧计算
        frame_list = []
        frame_count = -1
        # 重新获取视频路径，复位游标
        vid = cv2.VideoCapture(line_data)
        while True:
            ret, frame = vid.read()
            # 判断视频流中的帧是否存在
            if frame is None:
                break
            frame_count += 1
            if frame_count not in frame_span:
                continue
            frame = cv2.resize(frame, (size[1], size[2]))
            frame_list.append(frame)
        data_list.append(frame_list)
        label_list.append(line_label)
    return np.asarray(data_list, np.float32), np.asarray(label_list, np.int32), len(classes)


if __name__ == '__main__':
    split_data(DATASET_DIR, OUTPUT_LIST)
    print('Processing Completed!')
